package com.ittao.taoaicodemother.utils;

import cn.hutool.core.img.ImgUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IORuntimeException;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import com.ittao.taoaicodemother.exception.BusinessException;
import com.ittao.taoaicodemother.exception.ErrorCode;
import io.github.bonigarcia.wdm.WebDriverManager;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.openqa.selenium.JavascriptExecutor;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeOptions;
import org.openqa.selenium.support.ui.WebDriverWait;

import java.io.File;
import java.time.Duration;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 截图工具类 - 线程安全版本
 */
@Slf4j
public class WebScreenshotUtils {

    private static final int DEFAULT_WIDTH = 1600;
    private static final int DEFAULT_HEIGHT = 1200;
    private static final int MAX_DRIVER_COUNT = 5; // 最大WebDriver实例数

    // WebDriver对象池
    private static final ConcurrentLinkedQueue<WebDriver> driverPool = new ConcurrentLinkedQueue<>();
    private static final AtomicInteger currentDriverCount = new AtomicInteger(0);

    // ThreadLocal存储当前线程的WebDriver
    private static final ThreadLocal<WebDriver> threadLocalDriver = new ThreadLocal<>();

    // 静态初始化WebDriverManager
    static {
        try {
            WebDriverManager.chromedriver().setup();
            log.info("WebDriverManager初始化成功");
        } catch (Exception e) {
            log.error("WebDriverManager初始化失败", e);
        }
    }

    /**
     * 生成网页截图 - 线程安全版本
     * @param webUrl 要截图的网址
     * @return 截图文件路径
     */
    public static String saveWebPageScreenshot(String webUrl) {
        WebDriver driver = null;
        try {
            // 非空校验
            if (StrUtil.isBlank(webUrl)) {
                log.error("网页截图失败，url为空");
                return null;
            }

            // 获取WebDriver实例
            driver = getWebDriver();

            // 创建临时目录
            String rootPath = System.getProperty("user.dir") + "/tmp/screenshots/" + UUID.randomUUID().toString().substring(0, 8);
            FileUtil.mkdir(rootPath);

            // 图片后缀
            final String IMAGE_SUFFIX = ".png";
            // 原始图片路径
            String imageSavePath = rootPath + File.separator + RandomUtil.randomNumbers(5) + IMAGE_SUFFIX;

            // 访问网页
            driver.get(webUrl);
            // 等待页面加载
            waitForPageLoad(driver);

            // 截图
            byte[] screenshotBytes = ((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES);
            // 保存原始图片
            saveImage(screenshotBytes, imageSavePath);
            log.info("原始截图保存成功：{}", imageSavePath);

            // 压缩图片
            final String COMPRESSED_SUFFIX = "_compressed.jpg";
            String compressedImagePath = rootPath + File.separator + RandomUtil.randomNumbers(5) + COMPRESSED_SUFFIX;
            compressImage(imageSavePath, compressedImagePath);
            log.info("压缩截图保存成功：{}", compressedImagePath);

            // 删除原始图片
            FileUtil.del(imageSavePath);
            return compressedImagePath;

        } catch (Exception e) {
            log.error("网页截图失败，url: {}", webUrl, e);
            return null;
        } finally {
            // 释放WebDriver回到池中
            releaseWebDriver(driver);
        }
    }

    /**
     * 获取WebDriver实例（线程安全）
     */
    private static WebDriver getWebDriver() {
        // 先尝试从ThreadLocal获取
        WebDriver driver = threadLocalDriver.get();
        if (driver != null) {
            return driver;
        }

        // 尝试从对象池获取
        driver = driverPool.poll();
        if (driver != null) {
            threadLocalDriver.set(driver);
            return driver;
        }

        // 如果池中没有且未达到最大数量，创建新的
        if (currentDriverCount.get() < MAX_DRIVER_COUNT) {
            driver = createWebDriver();
            if (driver != null) {
                currentDriverCount.incrementAndGet();
                threadLocalDriver.set(driver);
                log.info("创建新的WebDriver实例，当前总数: {}", currentDriverCount.get());
                return driver;
            }
        }

        // 如果达到最大数量，等待一小段时间后重试从池中获取
        try {
            Thread.sleep(100); // 等待100ms
            driver = driverPool.poll();
            if (driver != null) {
                threadLocalDriver.set(driver);
                return driver;
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        // 最后兜底：创建临时实例（这个实例用完就销毁）
        log.warn("WebDriver池已满，创建临时实例");
        return createWebDriver();
    }

    /**
     * 释放WebDriver实例
     */
    private static void releaseWebDriver(WebDriver driver) {
        if (driver == null) {
            return;
        }

        try {
            // 清理当前页面，避免内存泄漏
            driver.manage().deleteAllCookies();

            // 从ThreadLocal中移除
            threadLocalDriver.remove();

            // 如果池未满，放回池中
            if (driverPool.size() < MAX_DRIVER_COUNT) {
                driverPool.offer(driver);
            } else {
                // 池已满，直接销毁
                driver.quit();
                currentDriverCount.decrementAndGet();
                log.info("销毁多余的WebDriver实例，当前总数: {}", currentDriverCount.get());
            }
        } catch (Exception e) {
            log.error("释放WebDriver时发生异常", e);
            try {
                driver.quit();
                currentDriverCount.decrementAndGet();
            } catch (Exception quitException) {
                log.error("强制销毁WebDriver失败", quitException);
            }
        }
    }

    /**
     * 创建WebDriver实例
     */
    private static WebDriver createWebDriver() {
        try {
            // 配置 Chrome 选项
            ChromeOptions options = new ChromeOptions();
            // 无头模式
            options.addArguments("--headless");
            // 禁用GPU（在某些环境下避免问题）
            options.addArguments("--disable-gpu");
            // 禁用沙盒模式（Docker环境需要）
            options.addArguments("--no-sandbox");
            // 禁用开发者shm使用
            options.addArguments("--disable-dev-shm-usage");
            // 设置窗口大小
            options.addArguments(String.format("--window-size=%d,%d", DEFAULT_WIDTH, DEFAULT_HEIGHT));
            // 禁用扩展
            options.addArguments("--disable-extensions");
            // 设置用户代理
            options.addArguments("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36");
            // 禁用图片加载以提高性能
            options.addArguments("--disable-images");
            // 禁用JavaScript（如果不影响截图效果）
            // options.addArguments("--disable-javascript");

            // 创建驱动
            WebDriver driver = new ChromeDriver(options);
            // 设置页面加载超时
            driver.manage().timeouts().pageLoadTimeout(Duration.ofSeconds(30));
            // 设置隐式等待
            driver.manage().timeouts().implicitlyWait(Duration.ofSeconds(10));

            return driver;
        } catch (Exception e) {
            log.error("创建WebDriver失败", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "创建WebDriver失败");
        }
    }

    /**
     * 应用关闭时清理所有WebDriver
     */
    @PreDestroy
    public static void destroyAll() {
        log.info("开始清理所有WebDriver实例");

        // 清理对象池中的实例
        WebDriver driver;
        while ((driver = driverPool.poll()) != null) {
            try {
                driver.quit();
            } catch (Exception e) {
                log.error("销毁WebDriver实例失败", e);
            }
        }

        // 清理ThreadLocal中的实例（通常在正常流程中已经清理）
        WebDriver threadDriver = threadLocalDriver.get();
        if (threadDriver != null) {
            try {
                threadLocalDriver.remove();
                threadDriver.quit();
            } catch (Exception e) {
                log.error("销毁ThreadLocal中的WebDriver失败", e);
            }
        }

        currentDriverCount.set(0);
        log.info("WebDriver清理完成");
    }

    /**
     * 保存图片到文件
     */
    private static void saveImage(byte[] imageBytes, String imagePath) {
        try {
            FileUtil.writeBytes(imageBytes, imagePath);
        } catch (Exception e) {
            log.error("保存图片失败: {}", imagePath, e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "保存图片失败");
        }
    }

    /**
     * 压缩图片
     * @param originImagePath 原始图片路径
     * @param compressedImagePath 压缩后图片路径
     */
    private static void compressImage(String originImagePath, String compressedImagePath) {
        // 压缩图片质量（0.1 = 10% 质量）
        final float COMPRESSION_QUALITY = 0.8f; // 压缩质量，范围0-1，值越小质量越低
        try {
            ImgUtil.compress(
                    FileUtil.file(originImagePath),
                    FileUtil.file(compressedImagePath),
                    COMPRESSION_QUALITY
            );
        } catch (IORuntimeException e) {
            log.error("压缩图片失败: {} -> {}", originImagePath, compressedImagePath, e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "压缩图片失败");
        }
    }

    /**
     * 等待页面加载完成
     * @param webDriver WebDriver 实例
     */
    private static void waitForPageLoad(WebDriver webDriver) {
        try {
            // 创建等待页面加载对象
            WebDriverWait wait = new WebDriverWait(webDriver, Duration.ofSeconds(10));
            // 等待 document.readyState 为 'complete'
            wait.until(driver -> Objects.equals(((JavascriptExecutor) driver)
                    .executeScript("return document.readyState"), "complete")
            );
            Thread.sleep(2000);
            log.info("页面加载完成");
        } catch (InterruptedException e) {
            log.error("等待页面加载时出现异常，继续执行截图", e);
        }
    }
}
